import os
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import torch
import torch.nn.functional as F
import torch.optim as optim
import matplotlib.pyplot as plt
from torch import nn
from torch import Tensor
from PIL import Image
import torchvision
from torchvision import transforms, models, datasets
from torchvision.transforms import Compose, Resize, ToTensor
import time
from sklearn.model_selection import train_test_split
import pickle
# import cv2

print(os.listdir("./input"))

# df=pd.read_pickle("./input/WM811K.pkl")
df=pd.read_pickle("./input/LSWMD.pkl")
# info()函数用于打印DataFrame的简要摘要，显示有关DataFrame的信息，包括索引的数据类型dtype和列的数据类型dtype，非空值的数量和内存使用情况
df.info()
print(df.info())
# head( )是指取数据的前n行数据，默认是前5行
df.head()
print(df.head())
# tail()方法一般用来对数据集进行矩阵形式的显示，默认显示为数据集的最后5行
df.tail()
df['failureNum']=df.failureType
df['trainTestNum']=df.trianTestLabel
mapping_type={'Center':0,'Donut':1,'Edge-Loc':2,'Edge-Ring':3,'Loc':4,'Random':5,'Scratch':6,'Near-full':7,'none':8}
mapping_traintest={'Training':0,'Test':1}
df=df.replace({'failureNum':mapping_type, 'trainTestNum':mapping_traintest})
# typelist = ['Center','Donut','Edge-Loc','Edge-Ring','Loc','Random','Scratch','Near-full','none']
# df_withlabel = df[(df["failureNum"] <= 8) & (df["failureNum"] >= 0)]
# df_withlabel =df_withlabel.reset_index()
df_withpattern = df[(df["failureNum"] <= 7) & (df["failureNum"] >= 0)]
df_withpattern = df_withpattern.reset_index()
df_nonpattern = df[(df["failureNum"] == 8)]

# df_Center = df[(df["failureNum"] == 0)]
# df_Dount = df[(df["failureNum"] == 1)]
# df_EdgeLoc = df[(df["failureNum"] == 2)]
# df_EdgeRing = df[(df["failureNum"] == 3)]
# df_Loc = df[(df["failureNum"] == 4)]
# df_Random = df[(df["failureNum"] == 5)]
# df_Scratch = df[(df["failureNum"] == 6)]
# df_Nearfull = df[(df["failureNum"] == 7)]


# x = df_withlabel["waferMap"]
# y = df_withlabel['failureNum']

filepath = "./input/train/"
for i in range(0, df_withpattern.shape[0]):
    index = df_withpattern.index[i]
    img = df_withpattern.waferMap[index]
    label = str(df_withpattern.failureType[index][0]).replace("['", "").replace("']","")
    traintestlabel = df_withpattern.trianTestLabel[index]
    if traintestlabel == 'Training':
        plt.imsave("./input/train/" + str(label) + "/" + str(label) + "_" + str(index) + ".jpg", img)
    else:
        plt.imsave("./input/val/" + str(label)+ "/" + str(label) + "_" + str(index) + ".jpg", img)
    # img.save("./input/data/" + str(label) + "/" + str(i) + ".jpg")

filepath = "./input/onpattern/"
for i in range(0, df_nonpattern.shape[0]):
    index = df_nonpattern.index[i]
    img = df_nonpattern.waferMap[index]
    label = str(df_nonpattern.failureType[index][0]).replace("['", "").replace("']","")
    traintestlabel = df_nonpattern.trianTestLabel[index]
    plt.imsave(filepath + str(index) + ".jpg", img)
filepath = "./input/nonlabel/"
for i in range(0, df.shape[0]):
    index = df.index[i]
    img = df.waferMap[index]
    type = df.failureType[index]
    if not df.failureType[index]:
        # label = str(df_nonlabel.failureType[index][0]).replace("['", "").replace("']","")
        plt.imsave(filepath + str(index) + ".jpg", img)


# x_train, x_test, y_train, y_test = train_test_split(x, y, test_size=0.2, random_state=42)


#
# train = df[df['trianTestLabel'] == "Training"]
# train = train.drop(['waferIndex'], axis = 1)
# test = df[df['trianTestLabel'] == "Test"]
# Center = df[df['failureType'] == "Center"]
# Donut = df[df['failureType'] == "Donut"]
# EdgeLoc = df[df['failureType'] == "Edge-Loc"]
# EdgeRing = df[df['failureType'] == "Edge-Ring"]
# Loc = df[df['failureType'] == "Loc"]
# Random = df[df['failureType'] == "Random"]
# Scratch = df[df['failureType'] == "Scratch"]
# Nearfull = df[df['failureType'] == "Near-full"]
# df = df.drop(['waferIndex'], axis = 1)
# Center_train = train[train['failureType'] == "Center"]
# Donut_train = train[train['failureType'] == "Donut"]
# EdgeLoc_train = train[train['failureType'] == "Edge-Loc"]
# EdgeRing_train = train[train['failureType'] == "Edge-Ring"]
# Loc_train = train[train['failureType'] == "Loc"]
# Random_train = train[train['failureType'] == "Random"]
# Scratch_train = train[train['failureType'] == "Scratch"]
# Nearfull_train = train[train['failureType'] == "Near-full"]
#
# Center_test = test[test['failureType'] == "Center"]
# Donut_test = test[test['failureType'] == "Donut"]
# EdgeLoc_test = test[test['failureType'] == "Edge-Loc"]
# EdgeRing_test = test[test['failureType'] == "Edge-Ring"]
# Loc_test = test[test['failureType'] == "Loc"]
# Random_test = test[test['failureType'] == "Random"]
# Scratch_test = test[test['failureType'] == "Scratch"]
# Nearfull_test = test[test['failureType'] == "Near-full"]
#
# print("train_total:" + str(train.shape[0]))
# print("test_total:" + str(test.shape[0]))
# print("Center_total:" + str(Center.shape[0]) + " " + "Center_train:"+ str(Center_train.shape[0]) + " " + "Center_test:" + str(Center_test.shape[0]))
# print("Donut_total:" + str(Donut.shape[0]) + " " + "Donut_train:"+ str(Donut_train.shape[0]) + " " + "Donut_test:" + str(Donut_test.shape[0]))
# print("EdgeLoc_total:" + str(EdgeLoc.shape[0]) + " " + "EdgeLoc_train:"+ str(EdgeLoc_train.shape[0]) + " " + "EdgeLoc_test:" + str(EdgeLoc_test.shape[0]))
# print("EdgeRing_total:" + str(EdgeRing.shape[0]) + " " + "EdgeRing_train:"+ str(EdgeRing_train.shape[0]) + " " + "EdgeRing_test:" + str(EdgeRing_test.shape[0]))
# print("Loc_total:" + str(Loc.shape[0]) + " " + "Loc_train:"+ str(Loc_train.shape[0]) + " " + "Loc_test:" + str(Loc_test.shape[0]))
# print("Random_total:" + str(Random.shape[0]) + " " + "Random_train:"+ str(Random_train.shape[0]) + " " + "Random_test:" + str(Random_test.shape[0]))
# print("Scratch_total:" + str(Scratch.shape[0]) + " " + "Scratch_train:"+ str(Scratch_train.shape[0]) + " " + "Scratch_test:" + str(Scratch_test.shape[0]))
# print("Nearfull_total:" + str(Nearfull.shape[0]) + " " + "Nearfull_train:"+ str(Nearfull_train.shape[0]) + " " + "Nearfull_test:" + str(Nearfull_test.shape[0]))
#
# print("total")
#
# data_transforms = {
#     'train':
#         transforms.Compose([
#         transforms.Resize(256, 256),
#         transforms.RandomRotation(45), #随机旋转， -45到45度之间随机选择
#         # transforms.CenterCrop(64), #从中心开始裁剪
#         transforms.RandomHorizontalFlip(p=0.5), #随机水平翻转，选择一个概率
#         transforms.RandomVerticalFlip(p=0.5), #随机垂直翻转
#         # transforms.ColorJitter(brightness=0.2, contrast=0.1, saturation=0.1, hue=0.1), #参数1为亮度， 参数2为对比度，参数3为饱和度，参数4为色相
#         # transforms.RandomGrayscale(p=0.024), #概率转化为灰度，3通道就是R=G=B
#         transforms.ToTensor(),
#         transforms.Normalize([0.485, 0.456, 0.406],[0.229, 0.224, 0.225])
#     ]),
#     'Vaild':
#         transforms.Compose([
#             transforms.Resize(256, 256),
#             transforms.ToTensor(),
#             transforms.Normalize([0.485, 0.456, 0.406],[0.229, 0.224, 0.225])
#         ])
# }
#
